/*
 * Copyright (C) 2025 ameise <ameise.wang@gmail.com> - All Rights Reserved
 *
 * This file is part of e3ds.
 *
 * e3ds is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * e3ds is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with e3ds. If not, see <https://www.gnu.org/licenses/>.
 */

package rbmap

import (
	"cmp"
	"log"
	"sync"
)

func newMap[K cmp.Ordered, V any]() RBMap[K, V] {
	m := RBMap[K, V]{leaf: &node[K, V]{color: black}}
	m.root = m.leaf
	m.headIter = &Iterator[K, V]{m: &m}
	m.endIter = &Iterator[K, V]{m: &m}
	m.headIter.next = m.endIter
	m.endIter.prev = m.headIter
	m.pool = newPool[K, V]()
	return m
}

type RBMap[K cmp.Ordered, V any] struct {
	leaf     *node[K, V]
	zeroK    K
	zeroV    V
	root     *node[K, V]
	count    int
	headIter *Iterator[K, V]
	endIter  *Iterator[K, V]
	pool     *sync.Pool
}

func (ts *RBMap[K, V]) Store(key K, val V) {
	_, _ = ts.loadAndStore(key, val)
}

func (ts *RBMap[K, V]) Load(key K) (val V, ok bool) {
	_, n := ts.search(key, nil, ts.root)
	if n != ts.leaf {
		return n.iter.ctx, true
	}
	return ts.zeroV, false
}

func (ts *RBMap[K, V]) LoadAndStore(key K, val V) (old V, ok bool) {
	return ts.loadAndStore(key, val)
}

func (ts *RBMap[K, V]) Delete(key K) (V, bool) {
	return ts.delete(key)
}

func (ts *RBMap[K, V]) Clean() map[K]V {
	if ts.count == 0 {
		return nil
	}
	out := make(map[K]V, ts.count)
	itr := ts.BeginIterator()
	for itr != ts.endIter {
		out[itr.KeyNoError()] = itr.ValueNoError()
		itr, _ = itr.Next()
	}
	ts.Reset()
	return out
}

func (ts *RBMap[K, V]) Reset() {
	ts.count = 0
	ts.root = ts.leaf
	ts.headIter.next = ts.endIter
	ts.endIter.prev = ts.headIter
}

func (ts *RBMap[K, V]) BeginIterator() *Iterator[K, V] {
	return ts.headIter.next
}

func (ts *RBMap[K, V]) EndIterator() *Iterator[K, V] {
	return ts.endIter
}

func (ts *RBMap[K, V]) Range(f func(*Iterator[K, V]) bool) {
	iter := ts.BeginIterator()
	for iter != ts.endIter {
		goNext := f(iter)
		if !goNext {
			break
		}
		iter, _ = iter.Next()
	}
}

func (ts *RBMap[K, V]) Search(key K, mode ...SearchMode) *Iterator[K, V] {
	if ts.root == ts.leaf {
		return ts.endIter
	}

	md := SearchModeET
	if len(mode) > 0 {
		if mode[0]&(SearchModeLT|SearchModeGT) == SearchModeLT|SearchModeGT {
			panic("rbmap search mode is invalid: cannot specify both LT and GT")
		}
		md = mode[0]
	}

	f, n := ts.search(key, nil, ts.root)

	if f == nil {
		f = n
	}

	if n == ts.leaf {
		if md == SearchModeET {
			return ts.endIter
		}
		if md&SearchModeLT != 0 {
			if f.iter.key < key {
				return f.iter
			} else {
				if f.iter.prev.n == nil {
					return ts.endIter
				}
				return f.iter.prev
			}
		}
		if f.iter.key > key {
			return f.iter
		} else {
			return f.iter.next
		}
	} else {
		if md&SearchModeET != 0 {
			return n.iter
		}
		if md == SearchModeLT {
			if n.iter.key < key {
				return f.iter
			} else {
				if n.iter.prev.n == nil {
					return ts.endIter
				}
				return n.iter.prev
			}
		}
		if n.iter.key > key {
			return n.iter
		} else {
			return n.iter.next
		}
	}
}

func (ts *RBMap[K, V]) Size() int {
	return ts.count
}

func (ts *RBMap[K, V]) linkIterator(n *node[K, V], f *node[K, V], key K) {
	if f != nil {
		if key < f.iter.key {
			f.left = n
			f.iter.prev.next = n.iter
			n.iter.prev = f.iter.prev
			n.iter.next = f.iter
			f.iter.prev = n.iter
		} else {
			f.right = n
			f.iter.next.prev = n.iter
			n.iter.next = f.iter.next
			n.iter.prev = f.iter
			f.iter.next = n.iter
		}
	} else {
		ts.root = n
		ts.headIter.next = n.iter
		n.iter.prev = ts.headIter
		n.iter.next = ts.endIter
		ts.endIter.prev = n.iter
	}
}

func (ts *RBMap[K, V]) loadAndStore(key K, val V) (old V, ok bool) {
	f, c := ts.search(key, nil, ts.root)
	if c == ts.leaf {
		n := newNode(ts, key, val)
		n.father = f
		ts.linkIterator(n, f, key)
		ts.insertCase1(n)
		ts.count++
		return ts.zeroV, false
	} else {
		old, c.iter.ctx = c.iter.ctx, val
		return old, true
	}
}

func (ts *RBMap[K, V]) search(key K, f, n *node[K, V]) (*node[K, V], *node[K, V]) {
	if n == ts.leaf {
		return f, n
	}

	if key == n.iter.key {
		return f, n
	} else if key < n.iter.key {
		return ts.search(key, n, n.left)
	} else {
		return ts.search(key, n, n.right)
	}
}

func (ts *RBMap[K, V]) delete(key K) (V, bool) {
	_, n := ts.search(key, nil, ts.root)
	if n == ts.leaf {
		return ts.zeroV, false
	}
	ctx := n.iter.ctx
	ts.deleteNode(n)
	ts.count--
	return ctx, true
}

func (ts *RBMap[K, V]) deleteNode(n *node[K, V]) (isSwap bool) {
	var ntd *node[K, V]

	if n.left != ts.leaf && n.right != ts.leaf {
		tmp := n.iter.next.n

		n.iter, tmp.iter = tmp.iter, n.iter
		n.iter.n, tmp.iter.n = n, tmp

		ntd = tmp
		isSwap = true
	} else {
		ntd = n
	}

	child := ntd.right
	if ntd.left != ts.leaf {
		child = ntd.left
	}

	if ntd.isRoot() {
		ts.root = child
		child.father = nil
		child.color = black
	} else {
		if ts.checkColor(ntd, black) {
			if ts.checkColor(child, red) {
				child.color = black
			} else {
				ntd.left = nil
				ntd.right = nil
				ts.deleteCase1(ntd)
			}
		}

		if ntd.father.left == ntd {
			ntd.father.left = child
		} else {
			ntd.father.right = child
		}

		if child != ts.leaf {
			child.father = ntd.father
		}
	}

	ntd.iter.prev.next = ntd.iter.next
	ntd.iter.next.prev = ntd.iter.prev

	freeNode(ntd)
	return
}

func (ts *RBMap[K, V]) leftRotate(n *node[K, V]) {
	g := n.grandfather()
	f := n.father
	c := n.left

	n.father = g
	if g == nil {
		ts.root = n
	} else {
		if g.left == f {
			g.left = n
		} else {
			g.right = n
		}
	}

	f.right = c
	if c != ts.leaf {
		c.father = f
	}

	n.left = f
	f.father = n
}

func (ts *RBMap[K, V]) rightRotate(n *node[K, V]) {
	g := n.grandfather()
	f := n.father
	c := n.right

	n.father = g
	if g == nil {
		ts.root = n
	} else {
		if g.left == f {
			g.left = n
		} else {
			g.right = n
		}
	}

	f.left = c
	if c != ts.leaf {
		c.father = f
	}

	n.right = f
	f.father = n
}

func (ts *RBMap[K, V]) insertCase1(n *node[K, V]) {
	if n.isRoot() {
		n.color = black
	} else {
		ts.insertCase2(n)
	}
}

func (ts *RBMap[K, V]) insertCase2(n *node[K, V]) {
	if !ts.checkColor(n.father, black) {
		ts.insertCase3(n)
	}
}

func (ts *RBMap[K, V]) insertCase3(n *node[K, V]) {
	if ts.checkColor(n.uncle(), red) {
		n.father.color = black
		n.grandfather().color = red
		n.uncle().color = black
		ts.insertCase1(n.grandfather())
	} else {
		ts.insertCase4(n)
	}
}

func (ts *RBMap[K, V]) insertCase4(n *node[K, V]) {
	if n == n.father.left {
		if n.father == n.grandfather().right {
			ts.rightRotate(n)
			n.color = black
			n.father.color = red
			ts.leftRotate(n)
		} else {
			n.father.color = black
			n.grandfather().color = red
			ts.rightRotate(n.father)
		}

	} else if n == n.father.right {
		if n.father == n.grandfather().left {
			ts.leftRotate(n)
			n.color = black
			n.father.color = red
			ts.rightRotate(n)
		} else {
			n.father.color = black
			n.grandfather().color = red
			ts.leftRotate(n.father)
		}
	}
}

func (ts *RBMap[K, V]) deleteCase1(n *node[K, V]) {
	if n.isRoot() {
		return
	}
	b := n.brother()
	if ts.checkColor(n.brother(), red) {
		n.father.color = red
		b.color = black
		if n == n.father.left {
			ts.leftRotate(b)
		} else {
			ts.rightRotate(b)
		}
	}
	ts.deleteCase2(n)
}

func (ts *RBMap[K, V]) deleteCase2(n *node[K, V]) {
	b := n.brother()
	if ts.checkColor(b.left, black) && ts.checkColor(b.right, black) {
		b.color = red
		if ts.checkColor(n.father, red) {
			n.father.color = black
		} else {
			ts.deleteCase1(n.father)
		}
	} else {
		ts.deleteCase3(n)
	}
}

func (ts *RBMap[K, V]) deleteCase3(n *node[K, V]) {
	b := n.brother()

	for {
		if (n == n.father.left && ts.checkColor(b.right, red)) || (n == n.father.right && ts.checkColor(b.left, red)) {
			b.color = n.father.color
			n.father.color = black
			if n == n.father.left {
				b.right.color = black
				ts.leftRotate(b)
			} else {
				b.left.color = black
				ts.rightRotate(b)
			}
			return
		} else {
			b.color = red
			if n == n.father.left {
				b.left.color = black
				ts.rightRotate(b.left)
			} else {
				b.right.color = black
				ts.leftRotate(b.right)
			}
			b = n.brother()
		}
	}
}

func (ts *RBMap[K, V]) checkBalance() {
	if ts.root == ts.leaf {
		log.Println("nil map")
		return
	}

	maxCnt := 0
	cnt := 0
	ts.checkBalance2(ts.root, &maxCnt, cnt)
	log.Println("check map ok，map high:", maxCnt)
}

func (ts *RBMap[K, V]) checkBalance2(n *node[K, V], maxCnt *int, cnt int) {
	if n.iter != nil {
		log.Printf("check node:%d   maxCnt:%d  cnt:%d\n", n.iter.key, *maxCnt, cnt)
	}
	if ts.checkColor(n, red) && (ts.checkColor(n.left, red) || ts.checkColor(n.right, red)) {
		panic("checkBalance failed")
	}
	if ts.checkColor(n, black) {
		cnt += 1
	}

	if n == ts.leaf {
		if *maxCnt == 0 {
			*maxCnt = cnt
		}
		if *maxCnt != cnt {
			panic("checkBalance failed")
		}
		return
	}

	if n.left != ts.leaf {
		ts.checkBalance2(n.left, maxCnt, cnt)
	} else {
		if *maxCnt == 0 {
			*maxCnt = cnt
		}
		if *maxCnt != cnt {
			panic("checkBalance failed")
		}
	}

	if n.right != ts.leaf {
		ts.checkBalance2(n.right, maxCnt, cnt)
	} else {
		if *maxCnt == 0 {
			*maxCnt = cnt
		}
		if *maxCnt != cnt {
			panic("checkBalance failed")
		}
	}
}

func (ts *RBMap[K, V]) checkColor(n *node[K, V], color int) bool {
	if n == ts.leaf {
		return black == color
	}
	return n.color == color
}
